#ifndef __RB_TREE_H__
#define __RB_TREE_H__

static const int RED = 0;
static const int BLACK = 1;

class node
{
public:
    node(int data);
    void setChild(node *);

    int data_;
    int color_;
    node *parent_, *lchild_, *rchild_;
};
int color(node *p)
{
    if (!p)
        return BLACK;
    return p->color_;
}

node::node(int data)
    : data_(data), color_(RED), parent_(nullptr), lchild_(nullptr), rchild_(nullptr)
{
}

void node::setChild(node *p)
{
    if (!p)
        return;
    if (p->data_ < data_)
        lchild_ = p;
    else if (p->data_ > data_)
        rchild_ = p;
}

class rb_tree
{
public:
    rb_tree();
    void rotate_left(node *);
    void rotate_right(node *);
    node *get_min(node *);
    node *get_max(node *);
    node *prev(node *);
    node *next(node *);

    void removeBalance(node *);

    node *insert(int);
    node *search(int, node *&);
    void remove(int);

    node *root_;
};

rb_tree::rb_tree()
    : root_(nullptr)
{
}

void rb_tree::rotate_left(node *n)
{
    node *p = n;
    node *q = p->rchild_;
    node *parent = n->parent_;

    if (nullptr == parent)
    {
        root_ = q;
    }
    else
    {
        if (parent->lchild_ == p)
            parent->lchild_ = q;
        else
            parent->rchild_ = q;
    }
    q->parent_ = parent;
    p->parent_ = q;

    p->rchild_ = q->lchild_;
    if (q->lchild_)
        q->lchild_->parent_ = p;
    q->lchild_ = p;
}

void rb_tree::rotate_right(node *n)
{
    node *p = n;
    node *q = p->lchild_;
    node *parent = n->parent_;

    if (nullptr == parent)
    {
        root_ = q;
    }
    else
    {
        if (parent->lchild_ == p)
            parent->lchild_ = q;
        else
            parent->rchild_ = q;
    }
    q->parent_ = parent;
    p->parent_ = q;

    p->lchild_ = q->rchild_;
    if (q->rchild_)
        q->rchild_->parent_ = p;
    q->rchild_ = p;
}

node *rb_tree::get_min(node *p)
{
    while (p->lchild_)
    {
        p = p->lchild_;
    }
    return p;
}
node *rb_tree::get_max(node *p)
{
    while (p->rchild_)
    {
        p = p->rchild_;
    }
    return p;
}
node *rb_tree::prev(node *p)
{
    if (p->lchild_)
        return get_max(p->lchild_);
    else
    {
        node *parent = p->parent_;
        while (parent && parent->lchild_ == p)
        {
            p = parent;
            parent = p->parent_;
        }
        return parent;
    }
}
node *rb_tree::next(node *p)
{
    if (p->rchild_)
        return get_min(p->rchild_);
    else
    {
        node *parent = p->parent_;
        while (parent && parent->rchild_ == p)
        {
            p = parent;
            parent = p->parent_;
        }
        return parent;
    }
}

node *rb_tree::search(int data, node *&parent)
{
    node *tmp = root_;
    while (tmp)
    {
        if (data == tmp->data_)
            return tmp;
        else
        {
            parent = tmp;
            if (data < tmp->data_)
                tmp = tmp->lchild_;
            else
                tmp = tmp->rchild_;
        }
    }
    return tmp;
}

node *rb_tree::insert(int data)
{
    node *parent = nullptr;
    node *insertNode = search(data, parent);
    if (insertNode)
        return insertNode;
    insertNode = new node(data);
    insertNode->parent_ = parent;

    if (nullptr == parent)
        root_ = insertNode;
    else
        parent->setChild(insertNode);

    while ((parent = insertNode->parent_) != nullptr && parent->color_ == RED)
    {
        node *grandpa = parent->parent_;
        if (parent == grandpa->lchild_)
        {
            node *uncle = grandpa->rchild_;
            if (uncle && RED == uncle->color_)
            {
                grandpa->color_ = RED;
                parent->color_ = BLACK;
                uncle->color_ = BLACK;
                insertNode = grandpa;
                continue;
            }
            else
            {
                if (insertNode == parent->rchild_)
                {
                    rotate_left(parent);
                    node *tmp = parent;
                    parent = insertNode;
                    insertNode = tmp;
                }
                parent->color_ = BLACK;
                grandpa->color_ = RED;
                rotate_right(grandpa);
            }
        }
        else if (parent == grandpa->rchild_)
        {
            node *uncle = grandpa->lchild_;
            if (uncle && RED == uncle->color_)
            {
                grandpa->color_ = RED;
                parent->color_ = BLACK;
                uncle->color_ = BLACK;
                insertNode = grandpa;
                continue;
            }
            else
            {
                if (insertNode == parent->lchild_)
                {
                    rotate_right(parent);
                    node *tmp = parent;
                    parent = insertNode;
                    insertNode = parent;
                }
                parent->color_ = BLACK;
                grandpa->color_ = RED;
                rotate_left(grandpa);
            }
        }
    }
    root_->color_ = BLACK;
    return insertNode;
}

void rb_tree::removeBalance(node *p)
{
    node *parent = p->parent_;
    node *sibling;
    while (BLACK == color(p) && parent)
    {
        if (parent->lchild_ == p)
        {
            sibling = parent->rchild_;
            if (RED == color(sibling))
            {
                sibling->color_ = BLACK;
                parent->color_ = RED;
                rotate_left(parent);
                sibling = parent->rchild_;
            }
            if (BLACK == color(sibling->lchild_) && BLACK == color(sibling->rchild_))
            {
                sibling->color_ = RED;
                p = parent;
                parent = p->parent_;
            }
            else
            {
                if (BLACK == color(sibling->rchild_))
                {
                    sibling->lchild_->color_ = BLACK;
                    sibling->color_ = RED;
                    rotate_right(sibling);
                    sibling = parent->rchild_;
                }
                sibling->color_ = parent->color_;
                parent->color_ = BLACK;
                sibling->rchild_->color_ = BLACK;
                rotate_left(parent);
                p = root_;
                break;
            }
        }
        else
        {
            sibling = parent->lchild_;
            if (RED == color(sibling))
            {
                sibling->color_ = BLACK;
                parent->color_ = RED;
                rotate_right(parent);
                sibling = parent->lchild_;
            }
            if (BLACK == color(sibling->lchild_) && BLACK == color(sibling->rchild_))
            {
                sibling->color_ = RED;
                p = parent;
                parent = p->parent_;
            }
            else
            {
                if (BLACK == color(sibling->lchild_))
                {
                    sibling->rchild_->color_ = BLACK;
                    sibling->color_ = RED;
                    rotate_left(sibling);
                    sibling = parent->lchild_;
                }
                sibling->color_ = parent->color_;
                parent->color_ = BLACK;
                sibling->lchild_->color_ = BLACK;
                rotate_right(parent);
                p = root_;
                break;
            }           
        }       
    }
    if (p)
        p->color_ = BLACK;
}
void rb_tree::remove(int data)
{
    node *parent;
    node *p = search(data, parent);
    if (nullptr == p)
        return;

    node *lchild = p->lchild_, *rchild = p->rchild_;
    if (lchild != nullptr && rchild != nullptr)
    {
        node *next = get_min(rchild);
        p->data_ = next->data_;
        p = next;
    }

    if (p->color_ == BLACK)
    {
        removeBalance(p);
    }

    node *child = p->rchild_ ? p->rchild_ : p->lchild_;
    if (!p->parent_)
    {
        root_ = child;
        if (root_)
            root_->color_ = BLACK;
    }
    else
    {
        if (p == p->parent_->lchild_)
            p->parent_->lchild_ = child;
        else
            p->parent_->rchild_ = child;
    }
    if (child)
        child->parent_ = p->parent_;
    delete p;
}

#endif